module operators_nh_mod
  use gptl_timer_mod
  use spark_cubed_sphere
  use params_mod
  use static_mod
  use state_mod
  use diag_mod
  use tend_mod
  use math_mod
  use physics_const_mod
  use phy_operator_mod
  use riemann_mod
  use ghost_interp_mod
  use hope_recon_mod
  use test_mod, only : case_omega

  implicit none

  private

  public operators_nh_run
  public operators_hor_run
  public operators_ver_run
  public Rayleigh_damping

contains

  subroutine operators_nh_run(state, tend, substep)
    type(state_type), intent(inout) :: state
    type(tend_type ), intent(inout) :: tend
    integer, intent(in) :: substep
    
    ret = gptlstart('Rayleigh_damping')
    call Rayleigh_damping(state,state_ref,dvs,dve)
    ret = gptlstop('Rayleigh_damping')
    
    ret = gptlstart('reconstruct')
    call reconstruct(state, diag)
    ret = gptlstop('reconstruct')
    
    ret = gptlstart('correct_no_flux_bdy')
    call correct_no_flux_bdy(state, diag)
    ret = gptlstop('correct_no_flux_bdy')
    
    ret = gptlstart('unify_domain_edges')
    call unify_domain_edges(state, diag)
    ret = gptlstop('unify_domain_edges')
    
    ret = gptlstart('calc_numerical_flux')
    call calc_numerical_flux(state, diag)
    ret = gptlstop('calc_numerical_flux')
    
    ret = gptlstart('zero_out_source_terms')
    call zero_out_source_terms(state, diag)
    ret = gptlstop('zero_out_source_terms')
    
    ret = gptlstart('calc_source_terms')
    call calc_source_terms(state, diag)
    ret = gptlstop('calc_source_terms')
    
    ret = gptlstart('calc_tendency')
    call calc_tendency(state, diag, tend)
    ret = gptlstop('calc_tendency')

  end subroutine operators_nh_run

  subroutine operators_hor_run(state, tend)
    type(state_type), intent(inout) :: state
    type(tend_type ), intent(inout) :: tend
    
    ret = gptlstart('Rayleigh_damping')
    call Rayleigh_damping(state,state_ref,dvs,dve)
    ret = gptlstop('Rayleigh_damping')
    
    ret = gptlstart('reconstruct_hor')
    call reconstruct_hor(state, diag)
    ret = gptlstop('reconstruct_hor')
    
    !ret = gptlstart('deconvolution_hor')
    !call deconvolution_hor(state, diag)
    !ret = gptlstop('deconvolution_hor')
    
    ret = gptlstart('unify_domain_edges')
    call unify_domain_edges(state, diag)
    ret = gptlstop('unify_domain_edges')
    
    ret = gptlstart('calc_numerical_flux_hor')
    call calc_numerical_flux_hor(state, diag)
    ret = gptlstop('calc_numerical_flux_hor')
    
    ret = gptlstart('zero_out_source_terms_hor')
    call zero_out_source_terms_hor(state, diag)
    ret = gptlstop('zero_out_source_terms_hor')
    
    ret = gptlstart('calc_source_terms_hor')
    call calc_source_terms_hor(state, diag)
    ret = gptlstop('calc_source_terms_hor')
    
    ret = gptlstart('calc_tendency_hor')
    call calc_tendency_hor(state, diag, tend)
    ret = gptlstop('calc_tendency_hor')

  end subroutine operators_hor_run

  subroutine operators_ver_run(state, tend, idom, i, j)
    type(state_type), intent(inout) :: state
    type(tend_type ), intent(inout) :: tend
    integer         , intent(in   ) :: idom, i, j
    
    ret = gptlstart('reconstruct_ver')
    call reconstruct_ver(state, diag, idom, i, j)
    ret = gptlstop('reconstruct_ver')
    
    ret = gptlstart('correct_no_flux_bdy_single_column')
    call correct_no_flux_bdy_single_column(state, diag, idom, i, j)
    ret = gptlstop('correct_no_flux_bdy_single_column')
    
    ret = gptlstart('calc_numerical_flux_ver')
    call calc_numerical_flux_ver(state, diag, idom, i, j)
    ret = gptlstop('calc_numerical_flux_ver')
    
    ret = gptlstart('zero_out_source_terms_ver')
    call zero_out_source_terms_ver(state, diag, idom, i, j)
    ret = gptlstop('zero_out_source_terms_ver')
    
    ret = gptlstart('calc_source_terms_ver')
    call calc_source_terms_ver(state, diag, idom, i, j)
    ret = gptlstop('calc_source_terms_ver')
    
    ret = gptlstart('calc_tendency_ver')
    call calc_tendency_ver(state, diag, tend, idom, i, j)
    ret = gptlstop('calc_tendency_ver')

  end subroutine operators_ver_run

  subroutine unify_domain_edges(state, diag)
    type(state_type), intent(inout) :: state
    type(diag_type ), intent(inout) :: diag

    integer idom, idom_ngb, orient, orient_adj
    integer i, j, k, p, i_ngb, j_ngb, p_ngb, i_adj, j_adj, p_adj
    integer ids, ide, jds, jde, kds, kde, pes(6), pee(6)
    integer is, ie, js, je, di, dj
    real(r8), pointer, dimension(:,:,:,:,:) :: q, q_ngb
    real(r8), dimension(2,2) :: iAh, Ah_ngb
    real(r8) rhod, uc, vc, us, vs, w, theta, mr
    real(r8) sqrtGrho
    
    associate (mesh => state%mesh)
    do idom = 1, mesh%ndom
      if (mesh%domain(idom)%active) then
        call mesh%get_params(idom=idom, ids=ids, ide=ide, jds=jds, jde=jde, kds=kds, kde=kde, pes=pes, pee=pee)
        do orient = 1, 4
          if (mesh%domain(idom)%touch_edges(orient)) then
            idom_ngb = mesh%domain(idom)%ngb(orient)%ptr%id
            associate (Jab     => mesh%domain(idom    )%Jab,  &
                       Jab_ngb => mesh%domain(idom_ngb)%Jab,  &
                       ngb     => mesh%domain(idom)%ngb(orient))
            is = ids; ie = ide; js = jds; je = jde; di = 0; dj = 0
            if (orient == 1) then; ie = ids; di = -1; orient_adj = 2; q => diag%domain(idom)%qR; end if
            if (orient == 2) then; is = ide; di = +1; orient_adj = 1; q => diag%domain(idom)%qL; end if
            if (orient == 3) then; je = jds; dj = -1; orient_adj = 4; q => diag%domain(idom)%qT; end if
            if (orient == 4) then; js = jde; dj = +1; orient_adj = 3; q => diag%domain(idom)%qB; end if
            if (ngb%orient == 1) q_ngb => diag%domain(idom_ngb)%qL
            if (ngb%orient == 2) q_ngb => diag%domain(idom_ngb)%qR
            if (ngb%orient == 3) q_ngb => diag%domain(idom_ngb)%qB
            if (ngb%orient == 4) q_ngb => diag%domain(idom_ngb)%qT
            !$OMP PARALLEL DO PRIVATE(j, i, p, i_ngb, j_ngb, p_ngb, i_adj, j_adj, p_adj, &
            !$OMP                     rhod, uc, vc, us, vs, w, theta, mr, sqrtGrho, Ah_ngb, iAh) COLLAPSE(4)
            do k = kds, kde
              do j = js, je
                do i = is, ie
                  do p = pes(orient), pee(orient)
                    p_adj = p - pes(orient) + pes(orient_adj)
                    i_adj = i + di
                    j_adj = j + dj
                    call mesh%translate_idx_along_edge(idom, orient, i, j, p, i_ngb, j_ngb, p_ngb)
                    
                    sqrtGrho = q_ngb(p_ngb,i_ngb,j_ngb,k,1) + q_ngb(p_ngb,i_ngb,j_ngb,k,6)
                    rhod     = q_ngb(p_ngb,i_ngb,j_ngb,k,1) / Jab_ngb(p_ngb,i_ngb,j_ngb,k)
                    uc       = q_ngb(p_ngb,i_ngb,j_ngb,k,2) / sqrtGrho
                    vc       = q_ngb(p_ngb,i_ngb,j_ngb,k,3) / sqrtGrho
                    w        = q_ngb(p_ngb,i_ngb,j_ngb,k,4) / sqrtGrho
                    theta    = q_ngb(p_ngb,i_ngb,j_ngb,k,5) / q_ngb(p_ngb,i_ngb,j_ngb,k,1)
                    mr       = q_ngb(p_ngb,i_ngb,j_ngb,k,6) / q_ngb(p_ngb,i_ngb,j_ngb,k,1)
                    
                    Ah_ngb = mesh%get_dom_Ah (idom_ngb,p_ngb,i_ngb,j_ngb,k)
                    iAh    = mesh%get_dom_iAh(idom    ,p_adj,i_adj,j_adj,k)
                    call contrav_to_spherev(Ah_ngb, uc, vc, us, vs)
                    call spherev_to_contrav(iAh   , us, vs, uc, vc)
                    
                    q(p_adj,i_adj,j_adj,k,1) = Jab(p_adj,i_adj,j_adj,k  ) * rhod
                    q(p_adj,i_adj,j_adj,k,5) = q  (p_adj,i_adj,j_adj,k,1) * theta
                    q(p_adj,i_adj,j_adj,k,6) = q  (p_adj,i_adj,j_adj,k,1) * mr
                    sqrtGrho = q(p_adj,i_adj,j_adj,k,1) + q(p_adj,i_adj,j_adj,k,6)
                    q(p_adj,i_adj,j_adj,k,2) = sqrtGrho * uc
                    q(p_adj,i_adj,j_adj,k,3) = sqrtGrho * vc
                    q(p_adj,i_adj,j_adj,k,4) = sqrtGrho * w
                  end do
                end do
              end do
            end do
            !$OMP END PARALLEL DO
            end associate
          end if
        end do
      end if
    end do
    end associate

  end subroutine unify_domain_edges

  subroutine calc_numerical_flux(state, diag)
    type(state_type), intent(inout) :: state
    type(diag_type ), intent(inout) :: diag
    
    real(r_kind) :: fxc(-1:nz+2), fyc(-1:nz+2) ! Use for convolution
    real(r_kind), dimension(2,2) :: iGhL, iGhR
    real(r_kind), dimension(3,3) :: iGL, iGR
    real(r_kind) :: JhL, JhR
    integer idom, ids, ide, jds, jde, kds, kde, pes(6), pee(6)
    integer i, j, k, p, p_adj, iVar
  
    associate (mesh => state%mesh)
    do idom = 1, mesh%ndom
      if (mesh%domain(idom)%active) then
        call mesh%get_params(idom=idom, pes=pes, pee=pee, ids=ids, ide=ide, jds=jds, jde=jde, kds=kds, kde=kde)
        associate (Jab => mesh %domain(idom)%Jab,                                &
                   iGv => mesh %domain(idom)%iGv, Jv  => mesh %domain(idom)%Jv , &
                   qL  => diag%domain(idom)%qL  , qR  => diag%domain(idom)%qR  , &
                   qB  => diag%domain(idom)%qB  , qT  => diag%domain(idom)%qT  , &
                   qD  => diag%domain(idom)%qD  , qU  => diag%domain(idom)%qU  , &
                   fxp => diag%domain(idom)%fxp , fyp => diag%domain(idom)%fyp , fzp => diag%domain(idom)%fzp,&
                   fx  => diag%domain(idom)%fx  , fy  => diag%domain(idom)%fy  , fz  => diag%domain(idom)%fz  )
        !$OMP PARALLEL
        ! Riemann solver
        !$OMP DO PRIVATE(p, p_adj, iVar, iGhL, iGhR) COLLAPSE(3)
        do k = kds, kde
          do j = jds, jde
            do i = ids, ide + 1
              do p = pes(1), pee(1)
                p_adj = p - pes(1) + pes(2)
                iGhL = mesh%get_dom_iGh(idom,p_adj,i-1,j,k)
                iGhR = mesh%get_dom_iGh(idom,p    ,i  ,j,k)
                fxp(p,i,j,k,:) = riemann_solver_hor(1, Jab(    p_adj,i-1,j,k  ), Jab(    p,i,j,k  ), &
                                                       iGhL                    , iGhR              , &
                                                       qR (    p_adj,i-1,j,k,:), qL(     p,i,j,k,:) )
              end do
            end do
          end do
        end do
        !$OMP END DO NOWAIT
        !$OMP DO PRIVATE(p, p_adj, iVar, iGhL, iGhR) COLLAPSE(3)
        do k = kds, kde
          do j = jds, jde + 1
            do i = ids, ide
              do p = pes(3), pee(3)
                p_adj = p - pes(3) + pes(4)
                iGhL = mesh%get_dom_iGh(idom,p_adj,i,j-1,k)
                iGhR = mesh%get_dom_iGh(idom,p    ,i,j  ,k)
                fyp(p,i,j,k,:) = riemann_solver_hor(2, Jab(    p_adj,i,j-1,k  ), Jab(    p,i,j,k  ), &
                                                       iGhL                    , iGhR              , &
                                                       qT (    p_adj,i,j-1,k,:), qB (    p,i,j,k,:) )
              end do
            end do
          end do
        end do
        !$OMP END DO NOWAIT
        !$OMP DO PRIVATE(p, p_adj, iVar, JhL, JhR, iGL, iGR) COLLAPSE(3)
        do k = kds, kde + 1
          do j = jds, jde
            do i = ids, ide
              do p = pes(5), pee(5)
                p_adj = p - pes(5) + pes(6)
                JhL = mesh%get_dom_Jh(idom,p_adj,i,j,k-1)
                JhR = mesh%get_dom_Jh(idom,p    ,i,j,k  )
                iGL = mesh%get_dom_iG(idom,p_adj,i,j,k-1)
                iGR = mesh%get_dom_iG(idom,p    ,i,j,k  )
                fzp(p,i,j,k,:) = riemann_solver_ver( Jab(    p_adj,i,j,k-1  ), Jab(    p,i,j,k),  &
                                                     JhL                     , JhR             ,  &
                                                     Jv (    p_adj,i,j,k-1  ), Jv (    p,i,j,k),  &
                                                     iGL                     , iGR             ,  &
                                                     iGv(:,:,p_adj,i,j,k-1  ), iGv(:,:,p,i,j,k),  &
                                                     qU (    p_adj,i,j,k-1,:), qD (    p,i,j,k,:) )
              end do
              do iVar = 1, nvar
                fz(i,j,k,iVar) = mesh%domain(idom)%edge_quad_ver(fzp(:,i,j,k,iVar))
              end do
            end do
          end do
        end do
        !$OMP END DO
        
        ! Convolution
        !$OMP DO COLLAPSE(3) PRIVATE(p,k,fxc)
        do iVar = 1, nVar
          do j = jds, jde
            do i = ids, ide + 1
              do p = pes(1), pee(1)
                fxc(1:nz) = fxp(p,i,j,kds:kde,iVar)
                do k = kds, kde
                  fxp(p,i,j,k,iVar) = convolution( fxc(k-2:k+2), k, kds, kde )
                enddo
              enddo
            enddo
          enddo
        enddo
        !$OMP END DO NOWAIT
        
        !$OMP DO COLLAPSE(3) PRIVATE(p,k,fyc)
        do iVar = 1, nVar
          do j = jds, jde + 1
            do i = ids, ide
              do p = pes(3), pee(3)
                fyc(1:nz) = fyp(p,i,j,kds:kde,iVar)
                do k = kds, kde
                  fyp(p,i,j,k,iVar) = convolution( fyc(k-2:k+2), k, kds, kde )
                enddo
              enddo
            enddo
          enddo
        enddo
        !$OMP END DO
        
        !$OMP DO COLLAPSE(4)
        do k = kds, kde
          do j = jds, jde
            do i = ids, ide + 1
              do iVar = 1, nvar
                fx(i,j,k,iVar) = mesh%domain(idom)%edge_quad(fxp(:,i,j,k,iVar))
              end do
            end do
          end do
        end do
        !$OMP END DO NOWAIT
        !$OMP DO COLLAPSE(4)
        do k = kds, kde
          do j = jds, jde + 1
            do i = ids, ide
              do iVar = 1, nvar
                fy(i,j,k,iVar) = mesh%domain(idom)%edge_quad(fyp(:,i,j,k,iVar))
              end do
            end do
          end do
        end do
        !$OMP END DO
        
        !$OMP END PARALLEL
        end associate
      end if
    end do
    end associate

  end subroutine calc_numerical_flux
  
  subroutine calc_numerical_flux_hor(state, diag)
    type(state_type), intent(inout) :: state
    type(diag_type ), intent(inout) :: diag
    
    real(r_kind), dimension(2,2) :: iGhL, iGhR
    real(r_kind) :: JhL, JhR
    integer idom, ids, ide, jds, jde, kds, kde, pes(6), pee(6)
    integer i, j, k, p, p_adj, iVar
  
    associate (mesh => state%mesh)
    do idom = 1, mesh%ndom
      if (mesh%domain(idom)%active) then
        call mesh%get_params(idom=idom, pes=pes, pee=pee, ids=ids, ide=ide, jds=jds, jde=jde, kds=kds, kde=kde)
        associate (Jab => mesh %domain(idom)%Jab,                                &
                   qL  => diag%domain(idom)%qL  , qR  => diag%domain(idom)%qR  , &
                   qB  => diag%domain(idom)%qB  , qT  => diag%domain(idom)%qT  , &
                   qD  => diag%domain(idom)%qD  , qU  => diag%domain(idom)%qU  , &
                   fxp => diag%domain(idom)%fxp , fyp => diag%domain(idom)%fyp , &
                   fx  => diag%domain(idom)%fx  , fy  => diag%domain(idom)%fy    )
        !$OMP PARALLEL
        ! Riemann solver
        !$OMP DO PRIVATE(p, p_adj, iVar, iGhL, iGhR) COLLAPSE(3)
        do k = kds, kde
          do j = jds, jde
            do i = ids, ide + 1
              do p = pes(1), pee(1)
                p_adj = p - pes(1) + pes(2)
                iGhL = mesh%get_dom_iGh(idom,p_adj,i-1,j,k)
                iGhR = mesh%get_dom_iGh(idom,p    ,i  ,j,k)
                fxp(p,i,j,k,:) = riemann_solver_hor(1, Jab(    p_adj,i-1,j,k  ), Jab(    p,i,j,k  ), &
                                                       iGhL                    , iGhR              , &
                                                       qR (    p_adj,i-1,j,k,:), qL(     p,i,j,k,:) )
              end do
              do iVar = 1, nVar
                fx(i,j,k,iVar) = mesh%domain(idom)%edge_quad(fxp(:,i,j,k,iVar))
              enddo
            end do
          end do
        end do
        !$OMP END DO NOWAIT
        !$OMP DO PRIVATE(p, p_adj, iVar, iGhL, iGhR) COLLAPSE(3)
        do k = kds, kde
          do j = jds, jde + 1
            do i = ids, ide
              do p = pes(3), pee(3)
                p_adj = p - pes(3) + pes(4)
                iGhL = mesh%get_dom_iGh(idom,p_adj,i,j-1,k)
                iGhR = mesh%get_dom_iGh(idom,p    ,i,j  ,k)
                fyp(p,i,j,k,:) = riemann_solver_hor(2, Jab(    p_adj,i,j-1,k  ), Jab(    p,i,j,k  ), &
                                                       iGhL                    , iGhR              , &
                                                       qT (    p_adj,i,j-1,k,:), qB (    p,i,j,k,:) )
              end do
              do iVar = 1, nVar
                fy(i,j,k,iVar) = mesh%domain(idom)%edge_quad(fyp(:,i,j,k,iVar))
              enddo
            end do
          end do
        end do
        !$OMP END DO
        !$OMP END PARALLEL
        end associate
      end if
    end do
    end associate
  
  end subroutine calc_numerical_flux_hor
  
  subroutine calc_numerical_flux_ver(state, diag, idom, i, j)
    type(state_type), intent(inout) :: state
    type(diag_type ), intent(inout) :: diag
    integer         , intent(in   ) :: idom, i, j
  
    real(r_kind), dimension(3,3) :: iGL, iGR
    real(r_kind) :: JhL, JhR
    integer ids, ide, jds, jde, kds, kde, pes(6), pee(6)
    integer k, p, p_adj, iVar
  
    associate (mesh => state%mesh)
      if (mesh%domain(idom)%active) then
        call mesh%get_params(idom=idom, pes=pes, pee=pee, kds=kds, kde=kde)
        associate (Jab=> mesh %domain(idom)%Jab,                               &
                   iGv => mesh %domain(idom)%iGv, Jv => mesh %domain(idom)%Jv, &
                   qD  => diag%domain(idom)%qD  , qU => diag%domain(idom)%qU,  &
                   fzp => diag%domain(idom)%fzp , fz => diag%domain(idom)%fz   )
        do k = kds, kde + 1
          do p = pes(5), pee(5)
            p_adj = p - pes(5) + pes(6)
            JhL = mesh%get_dom_Jh(idom,p_adj,i,j,k-1)
            JhR = mesh%get_dom_Jh(idom,p    ,i,j,k  )
            iGL = mesh%get_dom_iG(idom,p_adj,i,j,k-1)
            iGR = mesh%get_dom_iG(idom,p    ,i,j,k  )
            fzp(p,i,j,k,:) = riemann_solver_ver( Jab(    p_adj,i,j,k-1  ), Jab(    p,i,j,k),  &
                                                 JhL                     , JhR             ,  &
                                                 Jv (    p_adj,i,j,k-1  ), Jv (    p,i,j,k),  &
                                                 iGL                     , iGR             ,  &
                                                 iGv(:,:,p_adj,i,j,k-1  ), iGv(:,:,p,i,j,k),  &
                                                 qU (    p_adj,i,j,k-1,:), qD (    p,i,j,k,:) )
          end do
          do iVar = 1, nvar
            fz(i,j,k,iVar) = mesh%domain(idom)%edge_quad_ver(fzp(:,i,j,k,iVar))
          end do
        end do
        end associate
      end if
    end associate
  
  end subroutine calc_numerical_flux_ver

  subroutine calc_source_terms(state, diag)
    type(state_type), intent(inout) :: state
    type(diag_type ), intent(inout) :: diag
    
    real(r8), dimension(nvq,nVar) :: SM, FC, GF
    real(r8) :: sqrtGrho ! sqrtG * rho
    real(r8) :: rho      ! density
    real(r8) :: p        ! pressure
    real(r8) :: u(3)     ! velocity vector
    real(r8) :: AVSR(3)  ! Angle Velocity of Self-Rotation
    real(r8) :: f        ! Coriolis parameter
    real(r8) :: xx, yy, delta, omega
    real(r8) :: gm(3,3,3)
    real(r8) :: Jh
    real(r8) :: iGh(3,3)
    real(r8) :: CS(3,3,3)
    
    integer idom, ids, ide, jds, jde, kds, kde, pc, pqs, pqe
    integer ip, pq, i, j, k, iVar
    
    associate (mesh => state%mesh)
    do idom = 1, mesh%ndom
      if (mesh%domain(idom)%active) then
        call mesh%get_params(idom=idom, pqs=pqs, pqe=pqe, ids=ids, ide=ide, jds=jds, jde=jde, kds=kds, kde=kde)
        associate (q     => diag%domain(idom)%qQ,        &
                   src   => diag%domain(idom)%src,       &
                   g     => static%domain(idom)%gravityQ,&
                   x     => mesh%domain(idom)%x,         &
                   y     => mesh%domain(idom)%y,         &
                   lat   => mesh%domain(idom)%lat,       &
                   Jab   => mesh%domain(idom)%Jab,       &
                   r     => mesh%domain(idom)%r,         &
                   omg   => mesh%domain(idom)%omg        )
        !$OMP PARALLEL DO COLLAPSE(3) &
        !!$OMP PRIVATE( sqrtGrho, rho, u, p, ip, pq, iVar, SM, FC, GF, AVSR )
        !$OMP PRIVATE( sqrtGrho, Jh, iGh, CS, rho, u, p, ip, pq, iVar, SM, FC, GF, AVSR, xx, yy, delta, omega, f, gm )
        do k = kds, kde
          do j = jds, jde
            do i = ids, ide
              
              do ip = 1, nvq
                pq = pqs + ip - 1
                
                Jh           = mesh%get_dom_Jh (idom,pq,i,j,k)
                iGh          = 0
                iGh(1:2,1:2) = mesh%get_dom_iGh(idom,pq,i,j,k)
                iGh(3,3)     = 1
                CS           = mesh%get_dom_CS (idom,pq,i,j,k)
                
                sqrtGrho = q(pq,i,j,k,1) + q(pq,i,j,k,6)
                u(1)     = q(pq,i,j,k,2) / sqrtGrho
                u(2)     = q(pq,i,j,k,3) / sqrtGrho
                u(3)     = q(pq,i,j,k,4) / sqrtGrho
                rho      = sqrtGrho / Jab(pq,i,j,k)
                p        = calc_pressure( Jab(pq,i,j,k), q(pq,i,j,k,:) )
                AVSR     = cross( omg(:,pq,i,j,k), u )
                
                SM(ip,:) = 0
                FC(ip,:) = 0
                GF(ip,:) = 0
                
                !! General Scheme
                !! Calculate metric term
                !SM(ip,2:4) = - Jab(pq,i,j,k) * calc_metric_term( CS, rho, u, p, iGh )
                !
                !! Calculate Coriolis force term
                !FC(ip,2:4) = -2 * Jab(pq,i,j,k) * Jh * rho * matmul( AVSR, iGh )
                !
                !! Calculate Gravity
                !GF(ip,4) = -sqrtGrho * g(pq,i,j,k)
                
                !! Simplified Scheme
                !xx    = tan( x(pq,i,j) )
                !yy    = tan( y(pq,i,j) )
                !delta = sqrt( 1._r_kind + xx*xx + yy*yy )
                !omega = case_omega
                !
                !! Calculate metric term
                !SM(ip,2) = Jab(pq,i,j,k) * (2./delta/delta)*(-xx*yy*yy*rho*u(1)*u(1)+yy*(1+yy*yy)*rho*u(1)*u(2)-delta*delta/r(pq,i,j,k)*rho*u(1)*u(3))
                !SM(ip,3) = Jab(pq,i,j,k) * (2./delta/delta)*(-xx*xx*yy*rho*u(2)*u(2)+xx*(1+xx*xx)*rho*u(1)*u(2)-delta*delta/r(pq,i,j,k)*rho*u(2)*u(3))
                !SM(ip,4) = Jab(pq,i,j,k) * ( 2 * p / r(pq,i,j,k) &
                !                           + r(pq,i,j,k) * ( 1 + xx**2 ) * ( 1 + yy**2 ) / delta**4 &
                !                           * ( ( 1 + xx*2 ) * rho * u(1) * u(1) + ( 1 + yy**2 ) * rho * u(2) * u(2) - 2 * xx * yy * rho * u(1) * u(2) ) )
                !
                !! Calculate Coriolis force term
                !if(idom<=4)then
                !  FC(ip,2) = Jab(pq,i,j,k)*2.*Omega/delta/delta*(-xx*yy*yy*rho*u(1)+yy*(1.+yy*yy)*rho*u(2)-delta**2/r(pq,i,j,k)*rho*u(3))
                !  FC(ip,3) = Jab(pq,i,j,k)*2.*Omega/delta/delta*(-yy*(1.+xx*xx)*rho*u(1)+xx*yy*yy*rho*u(2)-delta**2*xx*yy/(r(pq,i,j,k)*(1+yy**2))*rho*u(3))
                !  FC(ip,4) = Jab(pq,i,j,k)*2.*Omega/delta/delta*r(pq,i,j,k)*(1+xx*xx)*rho*u(1)
                !elseif(idom==5)then
                !  FC(ip,2) = Jab(pq,i,j,k)*2.*Omega/delta/delta*(-xx*yy*rho*u(1)+(1.+yy*yy)*rho*u(2)+delta**2*yy/(r(pq,i,j,k)*(1+xx*xx))*rho*u(3))
                !  FC(ip,3) = Jab(pq,i,j,k)*2.*Omega/delta/delta*(-(1.+xx*xx)*rho*u(1)+xx*yy*rho*u(2)-delta**2*xx/(r(pq,i,j,k)*(1+yy*yy))*rho*u(3))
                !  FC(ip,4) = Jab(pq,i,j,k)*2.*Omega/delta/delta*r(pq,i,j,k)*(-yy*(1+xx*xx)*rho*u(1)+xx*(1+yy*yy)*rho*u(2))
                !else
                !  FC(ip,2) = Jab(pq,i,j,k)*(-2.*Omega/delta/delta)*(-xx*yy*rho*u(1)+(1.+yy*yy)*rho*u(2)+delta**2*yy/(r(pq,i,j,k)*(1+xx*xx))*rho*u(3))
                !  FC(ip,3) = Jab(pq,i,j,k)*(-2.*Omega/delta/delta)*(-(1.+xx*xx)*rho*u(1)+xx*yy*rho*u(2)-delta**2*xx/(r(pq,i,j,k)*(1+yy*yy))*rho*u(3))
                !  FC(ip,4) = Jab(pq,i,j,k)*2.*Omega/delta/delta*r(pq,i,j,k)*(yy*(1+xx*xx)*rho*u(1)-xx*(1+yy*yy)*rho*u(2))
                !endif
                !
                !! Calculate Gravity
                !GF(ip,4) = -sqrtGrho * g(pq,i,j,k)
                
                ! Scheme 3
                omega     = case_omega
                f         = 2 * omega * sin(lat(pq,i,j))
                gm        = 0
                gm(1,1,1) = CS(1,1,1)
                gm(1,2,1) = CS(1,2,1)
                gm(1,2,2) = CS(1,2,2)
                gm(2,2,2) = CS(2,2,2)
                
                SM(ip,2) = - Jab(pq,i,j,k) * (     gm(1,1,1) * ( rho * u(1) * u(1) + p * iGh(1,1) ) + 2 * gm(1,2,1) * ( rho * u(1) * u(2) + p * iGh(1,2) ) )
                SM(ip,3) = - Jab(pq,i,j,k) * ( 2 * gm(1,2,2) * ( rho * u(1) * u(2) + p * iGh(1,2) ) +     gm(2,2,2) * ( rho * u(2) * u(2) + p * iGh(2,2) ) )
                
                FC(ip,2) = - Jab(pq,i,j,k) * Jh * rho * f * ( u(1) * iGh(2,1) - u(2) * iGh(1,1) )
                FC(ip,3) = - Jab(pq,i,j,k) * Jh * rho * f * ( u(1) * iGh(2,2) - u(2) * iGh(1,2) )
                
                GF(ip,4) = -sqrtGrho * g(pq,i,j,k)
                
                !!if( idom==6 )then
                !  print*,idom,i,j,k
                !  !print*,( FC(ip,2:4) + 2 * Jab(pq,i,j,k) * Jh * rho * matmul( AVSR, iGh ) ) / ( -2 * Jab(pq,i,j,k) * Jh * rho * matmul( AVSR, iGh ) )
                !  !print*,FC(ip,2:4)
                !  !print*,-2 * Jab(pq,i,j,k) * Jh * rho * matmul( AVSR, iGh )
                !  !print*,''
                !  print*,( SM(ip,2:4) + Jab(pq,i,j,k) * calc_metric_term( CS, rho, u, p, iGh ) ) / ( - Jab(pq,i,j,k) * calc_metric_term( CS, rho, u, p, iGh ) )
                !  print*,SM(ip,2:4)
                !  print*,- Jab(pq,i,j,k) * calc_metric_term( CS, rho, u, p, iGh )
                !  print*,''
                !!endif
              end do
              
              do iVar = 1, nvar
                src(i,j,k,iVar) = src(i,j,k,iVar) + mesh%domain(idom)%cell_quad( SM(:,iVar) + FC(:,iVar) + GF(:,iVar)  )
              end do
            end do
          end do
        end do
        !$OMP END PARALLEL DO
        end associate
      end if
    end do
    end associate

  end subroutine calc_source_terms

  subroutine calc_source_terms_hor(state, diag)

    type(state_type), intent(inout) :: state
    type(diag_type ), intent(inout) :: diag
    
    real(r8), dimension(nvq,nVar) :: SM, FC
    real(r8) :: sqrtGrho ! sqrtG * rho
    real(r8) :: rho      ! density
    real(r8) :: p        ! pressure
    real(r8) :: u(3)     ! velocity vector
    real(r8) :: AVSR(3)  ! Angle Velocity of Self-Rotation
    real(r8) :: f        ! Coriolis parameter
    real(r8) :: xx, yy, delta, omega
    real(r8) :: gm(3,3,3)
    real(r8) :: Jh
    real(r8) :: iGh(3,3)
    real(r8) :: CS(3,3,3)
    
    integer idom, ids, ide, jds, jde, kds, kde, pc, pqs, pqe
    integer ip, pq, i, j, k, iVar
    
    associate (mesh => state%mesh)
    do idom = 1, mesh%ndom
      if (mesh%domain(idom)%active) then
        call mesh%get_params(idom=idom, pqs=pqs, pqe=pqe, ids=ids, ide=ide, jds=jds, jde=jde, kds=kds, kde=kde)
        associate (q     => diag%domain(idom)%qQ,        &
                   src   => diag%domain(idom)%srch,      &
                   g     => static%domain(idom)%gravityQ,&
                   x     => mesh%domain(idom)%x,         &
                   y     => mesh%domain(idom)%y,         &
                   lat   => mesh%domain(idom)%lat,       &
                   Jab   => mesh%domain(idom)%Jab,       &
                   r     => mesh%domain(idom)%r,         &
                   omg   => mesh%domain(idom)%omg        )
        !$OMP PARALLEL DO COLLAPSE(3) &
        !!$OMP PRIVATE( sqrtGrho, rho, u, p, ip, pq, iVar, SM, FC, AVSR )
        !$OMP PRIVATE( sqrtGrho, Jh, iGh, CS, rho, u, p, ip, pq, iVar, SM, FC, AVSR, xx, yy, delta, omega, f, gm )
        do k = kds, kde
          do j = jds, jde
            do i = ids, ide
              
              do ip = 1, nvq
                pq = pqs + ip - 1
                
                Jh           = mesh%get_dom_Jh (idom,pq,i,j,k)
                iGh          = 0
                iGh(1:2,1:2) = mesh%get_dom_iGh(idom,pq,i,j,k)
                iGh(3,3)     = 1
                CS           = mesh%get_dom_CS (idom,pq,i,j,k)
                
                sqrtGrho = q(pq,i,j,k,1) + q(pq,i,j,k,6)
                u(1)     = q(pq,i,j,k,2) / sqrtGrho
                u(2)     = q(pq,i,j,k,3) / sqrtGrho
                u(3)     = q(pq,i,j,k,4) / sqrtGrho
                rho      = sqrtGrho / Jab(pq,i,j,k)
                p        = calc_pressure( Jab(pq,i,j,k), q(pq,i,j,k,:) )
                AVSR     = cross( omg(:,pq,i,j,k), u )
                
                SM(ip,:) = 0
                FC(ip,:) = 0
                
                !! General Scheme
                !! Calculate metric term
                !SM(ip,2:4) = - Jab(pq,i,j,k) * calc_metric_term( CS, rho, u, p, iGh )
                !
                !! Calculate Coriolis force term
                !FC(ip,2:4) = -2 * Jab(pq,i,j,k) * Jh * rho * matmul( AVSR, iGh )
                
                !! Simplified Scheme 1
                !xx    = tan( x(pq,i,j) )
                !yy    = tan( y(pq,i,j) )
                !delta = sqrt( 1._r_kind + xx*xx + yy*yy )
                !omega = case_omega
                !
                !! Calculate metric term
                !SM(ip,2) = Jab(pq,i,j,k) * (2./delta/delta)*(-xx*yy*yy*rho*u(1)*u(1)+yy*(1+yy*yy)*rho*u(1)*u(2))
                !SM(ip,3) = Jab(pq,i,j,k) * (2./delta/delta)*(-xx*xx*yy*rho*u(2)*u(2)+xx*(1+xx*xx)*rho*u(1)*u(2))
                !
                !! Calculate Coriolis force term
                !if(idom<=4)then
                !  FC(ip,2) = Jab(pq,i,j,k)*2.*Omega/delta/delta*(-xx*yy*yy*rho*u(1)+yy*(1.+yy*yy)*rho*u(2))
                !  FC(ip,3) = Jab(pq,i,j,k)*2.*Omega/delta/delta*(-yy*(1.+xx*xx)*rho*u(1)+xx*yy*yy*rho*u(2))
                !elseif(idom==5)then
                !  FC(ip,2) = Jab(pq,i,j,k)*2.*Omega/delta/delta*(-xx*yy*rho*u(1)+(1.+yy*yy)*rho*u(2))
                !  FC(ip,3) = Jab(pq,i,j,k)*2.*Omega/delta/delta*(-(1.+xx*xx)*rho*u(1)+xx*yy*rho*u(2))
                !else
                !  FC(ip,2) = Jab(pq,i,j,k)*(-2.*Omega/delta/delta)*(-xx*yy*rho*u(1)+(1.+yy*yy)*rho*u(2))
                !  FC(ip,3) = Jab(pq,i,j,k)*(-2.*Omega/delta/delta)*(-(1.+xx*xx)*rho*u(1)+xx*yy*rho*u(2))
                !endif
                
                !! Simplified Scheme 2
                !xx    = tan( x(pq,i,j) )
                !yy    = tan( y(pq,i,j) )
                !delta = sqrt( 1._r_kind + xx*xx + yy*yy )
                !omega = case_omega
                !
                !! Calculate metric term
                !SM(ip,2) = Jab(pq,i,j,k) * (2./delta/delta)*(-xx*yy*yy*rho*u(1)*u(1)+yy*(1+yy*yy)*rho*u(1)*u(2)-delta*delta/r(pq,i,j,k)*rho*u(1)*u(3))
                !SM(ip,3) = Jab(pq,i,j,k) * (2./delta/delta)*(-xx*xx*yy*rho*u(2)*u(2)+xx*(1+xx*xx)*rho*u(1)*u(2)-delta*delta/r(pq,i,j,k)*rho*u(2)*u(3))
                !SM(ip,4) = Jab(pq,i,j,k) * ( 2 * p / r(pq,i,j,k) &
                !                           + r(pq,i,j,k) * ( 1 + xx**2 ) * ( 1 + yy**2 ) / delta**4 &
                !                           * ( ( 1 + xx*2 ) * rho * u(1) * u(1) + ( 1 + yy**2 ) * rho * u(2) * u(2) - 2 * xx * yy * rho * u(1) * u(2) ) )
                !
                !! Calculate Coriolis force term
                !if(idom<=4)then
                !  FC(ip,2) = Jab(pq,i,j,k)*2.*Omega/delta/delta*(-xx*yy*yy*rho*u(1)+yy*(1.+yy*yy)*rho*u(2)-delta**2/r(pq,i,j,k)*rho*u(3))
                !  FC(ip,3) = Jab(pq,i,j,k)*2.*Omega/delta/delta*(-yy*(1.+xx*xx)*rho*u(1)+xx*yy*yy*rho*u(2)-delta**2*xx*yy/(r(pq,i,j,k)*(1+yy**2))*rho*u(3))
                !  FC(ip,4) = Jab(pq,i,j,k)*2.*Omega/delta/delta*r(pq,i,j,k)*(1+xx*xx)*rho*u(1)
                !elseif(idom==5)then
                !  FC(ip,2) = Jab(pq,i,j,k)*2.*Omega/delta/delta*(-xx*yy*rho*u(1)+(1.+yy*yy)*rho*u(2)+delta**2*yy/(r(pq,i,j,k)*(1+xx*xx))*rho*u(3))
                !  FC(ip,3) = Jab(pq,i,j,k)*2.*Omega/delta/delta*(-(1.+xx*xx)*rho*u(1)+xx*yy*rho*u(2)-delta**2*xx/(r(pq,i,j,k)*(1+yy*yy))*rho*u(3))
                !  FC(ip,4) = Jab(pq,i,j,k)*2.*Omega/delta/delta*r(pq,i,j,k)*(-yy*(1+xx*xx)*rho*u(1)+xx*(1+yy*yy)*rho*u(2))
                !else
                !  FC(ip,2) = Jab(pq,i,j,k)*(-2.*Omega/delta/delta)*(-xx*yy*rho*u(1)+(1.+yy*yy)*rho*u(2)+delta**2*yy/(r(pq,i,j,k)*(1+xx*xx))*rho*u(3))
                !  FC(ip,3) = Jab(pq,i,j,k)*(-2.*Omega/delta/delta)*(-(1.+xx*xx)*rho*u(1)+xx*yy*rho*u(2)-delta**2*xx/(r(pq,i,j,k)*(1+yy*yy))*rho*u(3))
                !  FC(ip,4) = Jab(pq,i,j,k)*2.*Omega/delta/delta*r(pq,i,j,k)*(yy*(1+xx*xx)*rho*u(1)-xx*(1+yy*yy)*rho*u(2))
                !endif
                
                ! Scheme 3
                omega     = case_omega
                f         = 2 * omega * sin(lat(pq,i,j))
                gm        = 0
                gm(1,1,1) = CS(1,1,1)
                gm(1,2,1) = CS(1,2,1)
                gm(1,2,2) = CS(1,2,2)
                gm(2,2,2) = CS(2,2,2)
                
                SM(ip,2) = - Jab(pq,i,j,k) * (     gm(1,1,1) * ( rho * u(1) * u(1) + p * iGh(1,1) ) + 2 * gm(1,2,1) * ( rho * u(1) * u(2) + p * iGh(1,2) ) )
                SM(ip,3) = - Jab(pq,i,j,k) * ( 2 * gm(1,2,2) * ( rho * u(1) * u(2) + p * iGh(1,2) ) +     gm(2,2,2) * ( rho * u(2) * u(2) + p * iGh(2,2) ) )
                
                FC(ip,2) = - Jab(pq,i,j,k) * Jh * rho * f * ( u(1) * iGh(2,1) - u(2) * iGh(1,1) )
                FC(ip,3) = - Jab(pq,i,j,k) * Jh * rho * f * ( u(1) * iGh(2,2) - u(2) * iGh(1,2) )
                
                !!if( idom==6 )then
                !  print*,idom,i,j,k
                !  !print*,( FC(ip,2:4) + 2 * Jab(pq,i,j,k) * Jh * rho * matmul( AVSR, iGh ) ) / ( -2 * Jab(pq,i,j,k) * Jh * rho * matmul( AVSR, iGh ) )
                !  !print*,FC(ip,2:4)
                !  !print*,-2 * Jab(pq,i,j,k) * Jh * rho * matmul( AVSR, iGh )
                !  !print*,''
                !  print*,( SM(ip,2:4) + Jab(pq,i,j,k) * calc_metric_term( CS, rho, u, p, iGh ) ) / ( - Jab(pq,i,j,k) * calc_metric_term( CS, rho, u, p, iGh ) )
                !  print*,SM(ip,2:4)
                !  print*,- Jab(pq,i,j,k) * calc_metric_term( CS, rho, u, p, iGh )
                !  print*,''
                !!endif
              end do
              
              do iVar = 1, nvar
                src(i,j,k,iVar) = src(i,j,k,iVar) + mesh%domain(idom)%cell_quad( SM(:,iVar) + FC(:,iVar) )
              end do
            end do
          end do
        end do
        !$OMP END PARALLEL DO
        end associate
      end if
    end do
    end associate

  end subroutine calc_source_terms_hor
  
  subroutine calc_source_terms_ver(state, diag, idom, i, j)
    type(state_type), intent(inout) :: state
    type(diag_type ), intent(inout) :: diag
    integer         , intent(in   ) :: idom, i, j
    
    real(r8), dimension(nVar) :: GF
    real(r8) :: sqrtGrho ! sqrtG * rho
    
    integer kds, kde, pc, pqs, pqe
    integer ip, pq, k, iVar
    
    associate (mesh => state%mesh)
      if (mesh%domain(idom)%active) then
        call mesh%get_params(idom=idom, pc=pc, pqs=pqs, pqe=pqe, kds=kds, kde=kde)
        associate (q   => state%domain(idom)%q,       &
                   src => diag%domain(idom)%srcv,     &
                   Jab => mesh%domain(idom)%Jab,      &
                   g   => static%domain(idom)%gravity )
        do k = kds, kde
          sqrtGrho = Jab(pc,i,j,k) * ( q(i,j,k,1) + q(i,j,k,6) )
          
          GF = 0
          
          ! Calculate Gravity
          GF(4) = -sqrtGrho * g(i,j,k)
          
          do iVar = 1, nVar
            src(i,j,k,iVar) = src(i,j,k,iVar) + GF(iVar)
          end do
        end do
        end associate
      end if
    end associate
    
  end subroutine calc_source_terms_ver
  
  subroutine zero_out_source_terms(state, diag)
    type(state_type), intent(inout) :: state
    type(diag_type ), intent(inout) :: diag
    
    integer idom, ids, ide, jds, jde, kds, kde
    integer i, j, k, iVar
    
    associate (mesh => diag%mesh)
    do idom = 1, mesh%ndom
      if (mesh%domain(idom)%active) then
        call mesh%get_params(idom=idom, ids=ids, ide=ide, jds=jds, jde=jde, kds=kds, kde=kde)
        associate (src => diag%domain(idom)%src)
        !$OMP PARALLEL DO COLLAPSE(4)
        do iVar = 1, nVar
          do k = kds, kde
            do j = jds, jde
              do i = ids, ide
                src (i,j,k,iVar) = 0
              end do
            end do
          end do
        end do
        !$OMP END PARALLEL DO
        end associate
      end if
    end do
    end associate
  end subroutine zero_out_source_terms
  
  subroutine zero_out_source_terms_hor(state, diag)
    type(state_type), intent(inout) :: state
    type(diag_type ), intent(inout) :: diag
    
    integer idom, ids, ide, jds, jde, kds, kde
    integer i, j, k, iVar
    
    associate (mesh => diag%mesh)
    do idom = 1, mesh%ndom
      if (mesh%domain(idom)%active) then
        call mesh%get_params(idom=idom, ids=ids, ide=ide, jds=jds, jde=jde, kds=kds, kde=kde)
        associate (src => diag%domain(idom)%srch)
        !$OMP PARALLEL DO COLLAPSE(4)
        do iVar = 1, nVar
          do k = kds, kde
            do j = jds, jde
              do i = ids, ide
                src(i,j,k,iVar) = 0
              end do
            end do
          end do
        end do
        !$OMP END PARALLEL DO
        end associate
      end if
    end do
    end associate
  end subroutine zero_out_source_terms_hor
  
  subroutine zero_out_source_terms_ver(state, diag, idom, i, j)
    type(state_type), intent(inout) :: state
    type(diag_type ), intent(inout) :: diag
    integer         , intent(in   ) :: idom, i, j
    
    integer ids, ide, jds, jde, kds, kde
    integer k, iVar
    
    associate (mesh => diag%mesh)
      if (mesh%domain(idom)%active) then
        call mesh%get_params(idom=idom, kds=kds, kde=kde)
        associate (src => diag%domain(idom)%srcv)
        do iVar = 1, nVar
          do k = kds, kde
             src (i,j,k,iVar) = 0
          end do
        end do
        end associate
      end if
    end associate
  end subroutine zero_out_source_terms_ver
  
  subroutine calc_tendency(state, diag, tend)
    type(state_type), intent(inout) :: state
    type(diag_type ), intent(inout) :: diag
    type(tend_type ), intent(inout) :: tend

    integer idom, ids, ide, jds, jde, kds, kde, pqs, pqe
    integer i, j, k, iVar
    
    integer ims, ime, jms, jme, kms, kme

    associate (mesh => diag%mesh)
    do idom = 1, mesh%ndom
      if (mesh%domain(idom)%active) then
        call mesh%get_params(idom=idom, ids=ids, ide=ide, jds=jds, jde=jde, kds=kds, kde=kde, pqs=pqs, pqe=pqe)
        associate(dqdt => tend%domain(idom)%q,    &
                  src  => diag%domain(idom)%src,  &
                  fx   => diag%domain(idom)%fx,   &
                  fy   => diag%domain(idom)%fy,   &
                  fz   => diag%domain(idom)%fz,   &
                  dx   => mesh%domain(idom)%dx,   &
                  dy   => mesh%domain(idom)%dy,   &
                  deta => mesh%domain(idom)%deta  )
        !$OMP PARALLEL DO COLLAPSE(4)
        do iVar = 1, nvar
          do k = kds, kde
            do j = jds, jde
              do i = ids, ide
                dqdt(i,j,k,iVar) = - (fx(i+1,j,k,iVar) - fx(i,j,k,iVar)) / dx   &
                                   - (fy(i,j+1,k,iVar) - fy(i,j,k,iVar)) / dy   &
                                   - (fz(i,j,k+1,iVar) - fz(i,j,k,iVar)) / deta &
                                   + src(i,j,k,iVar)
                
                !!if(abs(static%domain(idom)%dzsdx(i,j))>30000.or.abs(static%domain(idom)%dzsdy(i,j))>30000)then
                !if(idom==6.and.iVar==4.and.i==23.and.j==23.and.k==38)then
                !  print*,idom,i,j,k,iVar
                !  print*,state%domain(idom)%q(i,j,k,iVar)
                !  print*,dqdt(i,j,k,iVar), state%domain(idom)%q(i,j,k,1), dqdt(i,j,k,iVar) / state%domain(idom)%q(i,j,k,1)
                !  print*,- (fx(i+1,j,k,iVar) - fx(i,j,k,iVar)) / dx,   fx(i,j,k,iVar), fx(i+1,j,k,iVar)
                !  print*,- (fy(i,j+1,k,iVar) - fy(i,j,k,iVar)) / dy,   fy(i,j,k,iVar), fy(i,j+1,k,iVar)
                !  print*,- (fz(i,j,k+1,iVar) - fz(i,j,k,iVar)) / deta, fz(i,j,k,iVar), fz(i,j,k+1,iVar)
                !  print*,src(i,j,k,iVar)
                !  print*,static%domain(idom)%dzsdx(i,j),static%domain(idom)%dzsdy(i,j)
                !endif
                
              end do
            end do
          end do
        end do
        !$OMP END PARALLEL DO
        end associate
      end if
    end do
    end associate
    
    !open(1,file='check_halo.txt')
    !idom = 1
    !k = kds
    !call state%mesh%get_params(idom=idom, ims=ims, ime=ime, jms=jms, jme=jme, kms=kms, kme=kme)
    !print*,ims,ime
    !do iVar = 1, nVar
    !  do j = jms,jme
    !    write(1,'(1x,49f30.5)')(state%domain(idom)%q(i,j,k,iVar),i=ims,ime)
    !  enddo
    !enddo
    !close(1)
    !stop 'Check calc_tendency'

  end subroutine calc_tendency
  
  subroutine calc_tendency_hor(state, diag, tend)
    type(state_type), intent(inout) :: state
    type(diag_type ), intent(inout) :: diag
    type(tend_type ), intent(inout) :: tend

    integer idom, ids, ide, jds, jde, kds, kde, pc
    integer i, j, k, iVar

    associate (mesh => diag%mesh)
    do idom = 1, mesh%ndom
      if (mesh%domain(idom)%active) then
        call mesh%get_params(idom=idom, ids=ids, ide=ide, jds=jds, jde=jde, kds=kds, kde=kde, pc=pc)
        associate(dqdt => tend%domain(idom)%q,     &
                  src  => diag%domain(idom)%srch,  &
                  fx   => diag%domain(idom)%fx,    &
                  fy   => diag%domain(idom)%fy,    &
                  Jab  => mesh%domain(idom)%Jab,   &
                  dx   => mesh%domain(idom)%dx,    &
                  dy   => mesh%domain(idom)%dy     )
        !$OMP PARALLEL DO COLLAPSE(4)
        do iVar = 1, nvar
          do k = kds, kde
            do j = jds, jde
              do i = ids, ide
                dqdt(i,j,k,iVar) = ( - (fx(i+1,j,k,iVar) - fx(i,j,k,iVar)) / dx   &
                                     - (fy(i,j+1,k,iVar) - fy(i,j,k,iVar)) / dy   &
                                     + src(i,j,k,iVar) ) / Jab(pc,i,j,k)
              end do
            end do
          end do
        end do
        !$OMP END PARALLEL DO
        end associate
      end if
    end do
    end associate

  end subroutine calc_tendency_hor
  
  subroutine calc_tendency_ver(state, diag, tend, idom, i, j)
    type(state_type), intent(inout) :: state
    type(diag_type ), intent(inout) :: diag
    type(tend_type ), intent(inout) :: tend
    integer         , intent(in   ) :: idom,i,j

    integer ids, ide, jds, jde, kds, kde, pc
    integer k, iVar

    associate (mesh => diag%mesh)
      if (mesh%domain(idom)%active) then
        call mesh%get_params(idom=idom, ids=ids, ide=ide, jds=jds, jde=jde, kds=kds, kde=kde, pc=pc)
        associate(dqdt => tend%domain(idom)%q,     &
                  src  => diag%domain(idom)%srcv,  &
                  fz   => diag%domain(idom)%fz,    &
                  Jab  => mesh%domain(idom)%Jab,   &
                  deta => mesh%domain(idom)%deta   )
        do iVar = 1, nvar
          do k = kds, kde
            dqdt(i,j,k,iVar) = ( - (fz(i,j,k+1,iVar) - fz(i,j,k,iVar)) / deta &
                                 + src(i,j,k,iVar) ) / Jab(pc,i,j,k)
          end do
        end do
        end associate
      end if
    end associate
  end subroutine calc_tendency_ver
  
  function calc_metric_term(CS, rho, u, p, iG) result(SM)
    real(r_kind) :: SM(3)
    real(r_kind) :: CS(3,3,3)
    real(r_kind) :: rho
    real(r_kind) :: u(3)
    real(r_kind) :: p
    real(r_kind) :: iG(3,3)
    
    integer i,j,k
    
    SM = 0
    do i = 1, 3
      do j = 1, 3
        do k = 1, 3
          SM(i) = SM(i) + CS(j,k,i) * ( rho * u(j) * u(k) + iG(j,k) * p )
        enddo
      enddo
    enddo
    
  end function calc_metric_term

  subroutine Rayleigh_damping(state,state_ref,vs,ve)
    type(state_type), intent(inout) :: state
    type(state_type), intent(in   ) :: state_ref
    integer(i_kind),  intent(in   ) :: vs,ve

    integer idom, ids, ide, jds, jde, kds, kde, pqs, pqe
    integer p, i, j, k, iVar

    associate (mesh => state%mesh)
    do idom = 1, mesh%ndom
      if (mesh%domain(idom)%active) then
        call mesh%get_params(idom=idom, ids=ids, ide=ide, jds=jds, jde=jde, kds=kds, kde=kde, pqs=pqs, pqe=pqe)
        associate(rdc   => static%domain(idom)%rdc , &
                  q     => state%domain(idom)%q    , &
                  q_ref => state_ref%domain(idom)%q  )
        !$OMP PARALLEL DO COLLAPSE(4)
        do iVar = vs,ve
          do k = kds, kde
            do j = jds, jde
              do i = ids, ide
                !!if(mesh%domain(idom)%z(1,i,j,k)>20000)then
                !  print*,rdc(i,j,k), q(i,j,k,iVar), q_ref(i,j,k,iVar), mesh%domain(idom)%z(1,i,j,k), maxval(mesh%domain(idom)%z(1,:,:,:))
                !!endif
                q(i,j,k,iVar) = q(i,j,k,iVar) - rdc(i,j,k) * ( q(i,j,k,iVar) - q_ref(i,j,k,iVar) )
              enddo
            enddo
          enddo
        end do
        !$OMP END PARALLEL DO
        end associate
      end if
    end do
    end associate
    
  end subroutine Rayleigh_damping
  
  ! Convert flux to no-flux boundary
  subroutine correct_no_flux_bdy(state, diag)
    type(state_type), intent(inout) :: state
    type(diag_type ), intent(inout) :: diag
    
    real(r8) :: mtx(3,3)
    real(r8) :: dzdx, dzdy, dzdlon, dzdlat
    real(r8) :: uc(3), us(3)
    real(r8), dimension(2,2) :: Ah, iAh
    integer idom, ids, ide, jds, jde, kds, kde, pes(6), pee(6)
    integer p, i, j, k, iVar

    associate (mesh => diag%mesh)
      do idom = 1, mesh%ndom
        if (mesh%domain(idom)%active) then
          call mesh%get_params(idom=idom, ids=ids, ide=ide, jds=jds, jde=jde, kds=kds, kde=kde, pes=pes, pee=pee)
          associate(dzsdxD => static%domain(idom)%dzsdxD, &
                    dzsdyD => static%domain(idom)%dzsdyD, &
                    qD     => diag%domain(idom)%qD,       &
                    qU     => diag%domain(idom)%qU        )
            !$OMP PARALLEL DO PRIVATE(i, k, p, Ah, iAh, mtx, uc, us, dzdx, dzdy, dzdlon, dzdlat) COLLAPSE(2)
            do j = jds, jde
              do i = ids, ide
                ! Bottom boundary
                k = kds
                do p = pes(5), pee(5)
                  
                  Ah  = mesh%get_dom_Ah (idom,p,i,j,k)
                  iAh = mesh%get_dom_iAh(idom,p,i,j,k)
            
                  uc = qD(p,i,j,k,2:4)
                  us(3) = uc(3)
                  call contrav_to_spherev(Ah, uc(1), uc(2), us(1), us(2))
                  
                  dzdx = dzsdxD(p,i,j)
                  dzdy = dzsdyD(p,i,j)
                  
                  call convert_hor_deriv_cube_to_sph(dzdlon,dzdlat,dzdx,dzdy,iAh)
                  mtx = calc_flux_convert_matrix(dzdlon, dzdlat)
                  us = matmul( mtx, us )
                  
                  uc(3) = us(3)
                  call spherev_to_contrav(iAh, us(1), us(2), uc(1), uc(2))
                  qD(p,i,j,k,2:4) = uc
                enddo
                qU(:,i,j,k-1,2:4) = qD(:,i,j,k,2:4) ! Fill up outside points for Riemann solver
                
                ! Top boundary
                k = kde
                do p = pes(6), pee(6)
                  qU(p,i,j,k,2:4) = 0
                enddo
                qD(:,i,j,k+1,2:4) = qU(:,i,j,k,2:4) ! Fill up outside points for Riemann solver
              enddo
            enddo
            !$OMP END PARALLEL DO
          end associate
        end if
      end do
    end associate
  end subroutine correct_no_flux_bdy
  
  ! Convert flux to no-flux boundary
  subroutine correct_no_flux_bdy_single_column(state, diag, idom, i, j)
    type(state_type), intent(inout) :: state
    type(diag_type ), intent(inout) :: diag
    integer         , intent(in   ) :: idom, i, j
    
    real(r8) :: mtx(3,3)
    real(r8) :: dzdx, dzdy, dzdlon, dzdlat
    real(r8) :: uc(3), us(3)
    real(r8), dimension(2,2) :: Ah, iAh
    integer ids, ide, jds, jde, kds, kde, pes(6), pee(6)
    integer p, k, iVar

    associate (mesh => diag%mesh)
      if (mesh%domain(idom)%active) then
        call mesh%get_params(idom=idom, ids=ids, ide=ide, jds=jds, jde=jde, kds=kds, kde=kde, pes=pes, pee=pee)
        associate(dzsdxD => static%domain(idom)%dzsdxD, &
                  dzsdyD => static%domain(idom)%dzsdyD, &
                  Jv     => mesh%domain(idom)%Jv,       &
                  iGv    => mesh%domain(idom)%iGv,      &
                  qD     => diag%domain(idom)%qD,       &
                  qU     => diag%domain(idom)%qU        )
          ! Bottom boundary
          k = kds
          do p = pes(5), pee(5)
            
            Ah  = mesh%get_dom_Ah (idom,p,i,j,k)
            iAh = mesh%get_dom_iAh(idom,p,i,j,k)
              
            uc = qD(p,i,j,k,2:4)
            us(3) = uc(3)
            call contrav_to_spherev(Ah, uc(1), uc(2), us(1), us(2))
            
            !dzdx = dzsdxD(p,i,j)
            !dzdy = dzsdyD(p,i,j)
            
            dzdx = -iGv(1,3,p,i,j,k) * Jv(p,i,j,k)
            dzdy = -iGv(2,3,p,i,j,k) * Jv(p,i,j,k)
            
            call convert_hor_deriv_cube_to_sph(dzdlon,dzdlat,dzdx,dzdy,iAh)
            mtx = calc_flux_convert_matrix(dzdlon, dzdlat)
            us = matmul( mtx, us )
            
            uc(3) = us(3)
            call spherev_to_contrav(iAh, us(1), us(2), uc(1), uc(2))
            qD(p,i,j,k,2:4) = uc
          enddo
          qU(:,i,j,k-1,2:4) = qD(:,i,j,k,2:4) ! Fill up outside points for Riemann solver
          
          ! Top boundary
          k = kde
          do p = pes(6), pee(6)
            qU(p,i,j,k,2:4) = 0
          enddo
          qD(:,i,j,k+1,2:4) = qU(:,i,j,k,2:4) ! Fill up outside points for Riemann solver
        end associate
      end if
    end associate
  end subroutine correct_no_flux_bdy_single_column
  
  ! Calculate convert matrix to convert flux for no-flux boundary condition
  function calc_flux_convert_matrix(dzdx, dzdy) result(mtx)
    real(r_kind) :: dzdx
    real(r_kind) :: dzdy
    real(r_kind) :: mtx(3,3)
    
    real(r_kind) :: n(3)
    
    n(1) = dzdx
    n(2) = dzdy
    n(3) = -1
    
    n = n / sqrt( dot_product( n, n ) )
    
    mtx(1,1) = 1._r_kind - n(1) * n(1)
    mtx(1,2) = -n(1) * n(2)
    mtx(1,3) = -n(1) * n(3)
    mtx(2,1) = mtx(1,2)
    mtx(2,2) = 1._r_kind - n(2) * n(2)
    mtx(2,3) = -n(2) * n(3)
    mtx(3,1) = mtx(1,3)
    mtx(3,2) = mtx(2,3)
    mtx(3,3) = 1._r_kind - n(3) * n(3)
  
  end function calc_flux_convert_matrix
  
end module operators_nh_mod
